import pygame
import random

# 初始化Pygame
pygame.init()

# 设置游戏窗口
screen_width = 600
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("五子棋")

# 定义颜色
board_color = (225, 195, 131)
player_color = (0, 0, 0)
computer_color = (255, 255, 255)

# 定义棋盘和棋子
board_size = 15
board_pos_x = 50
board_pos_y = 50
piece_size = 20
piece_space = 35
board = [[0] * board_size for i in range(board_size)]

# 定义游戏状态
game_over = False

# 定义判断胜利函数


def check_win(player):
    # 判断行
    for i in range(board_size):
        count = 0
        for j in range(board_size):
            if board[i][j] == player:
                count += 1
                if count == 5:
                    return True
            else:
                count = 0
    # 判断列
    for j in range(board_size):
        count = 0
        for i in range(board_size):
            if board[i][j] == player:
                count += 1
                if count == 5:
                    return True
            else:
                count = 0
    # 判断主对角线
    for i in range(board_size - 4):
        for j in range(board_size - 4):
            count = 0
            for k in range(5):
                if board[i+k][j+k] == player:
                    count += 1
                    if count == 5:
                        return True
                else:
                    count = 0
    # 判断副对角线
    for i in range(board_size - 4):
        for j in range(4, board_size):
            count = 0
            for k in range(5):
                if board[i+k][j-k] == player:
                    count += 1
                    if count == 5:
                        return True
                else:
                    count = 0
    return False

# 定义重绘函数


def redraw():
    # 绘制棋盘
    screen.fill(board_color)
    for i in range(board_size):
        pygame.draw.line(screen, (0, 0, 0), (board_pos_x, board_pos_y + i * piece_space),
                         (board_pos_x + (board_size - 1) * piece_space, board_pos_y + i * piece_space))
        pygame.draw.line(screen, (0, 0, 0), (board_pos_x + i * piece_space, board_pos_y),
                         (board_pos_x + i * piece_space, board_pos_y + (board_size - 1) * piece_space))
    # 绘制棋子
    for i in range(board_size):
        for j in range(board_size):
            if board[i][j] == 1:
                pygame.draw.circle(screen, player_color, (board_pos_x + j *
                                   piece_space, board_pos_y + i * piece_space), piece_size)
            elif board[i][j] == 2:
                pygame.draw.circle(screen, computer_color, (board_pos_x +
                                   j * piece_space, board_pos_y + i * piece_space), piece_size)
    # 更新屏幕
    pygame.display.update()

# 定义计算机下棋函数


def computer_drop_piece():
    while True:
        row = random.randint(0, board_size - 1)
        col = random.randint(0, board_size - 1)
        if board[row][col] == 0:
            board[row][col] = 2
            return


# 游戏循环
while not game_over:
    # 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game_over = True
        elif event.type == pygame.MOUSEBUTTONDOWN:
            # 获取鼠标点击位置
            x, y = pygame.mouse.get_pos()
            # 判断是否在棋盘内
            if x > board_pos_x and x < board_pos_x + (board_size - 1) * piece_space and y > board_pos_y and y < board_pos_y + (board_size - 1) * piece_space:
                # 判断该位置是否已有棋子
                row = int((y - board_pos_y) / piece_space)
                col = int((x - board_pos_x) / piece_space)
                if board[row][col] == 0:
                    # 玩家落子
                    board[row][col] = 1
                    # 判断是否胜利
                    if check_win(1):
                        print("恭喜你获胜！")
                        game_over = True
                    # 计算机落子
                    computer_drop_piece()
                    # 判断是否胜利
                    if check_win(2):
                        print("很遗憾你输了！")
                        game_over = True
                    # 重绘棋盘
                    redraw()

# 退出游戏
pygame.quit()
